/*-
 * Copyright (C) 2020, Shanghai Trusted Industry Control Platform Co., Ltd.
 * Created by 周正辉 (Zhenghui Zhou) <zhouzhenghui@ticpsh.com>
 */

#ifndef FIXED_ADDRESS_H
#define FIXED_ADDRESS_H

#include <clang/AST/PrettyPrinter.h>
#include <clang/AST/RecursiveASTVisitor.h>

#include <vector>

namespace rocket {
namespace share {

/**
 * @brief 绝对地址表达式的内部表示
 */
struct FixedAddress {
  const clang::Expr *parent;    /**< 不重复的父级指针表达式 */
  const clang::Expr *expr;      /**< 绝对地址表达式 */
  clang::QualType type;         /**< 指针类型 */
  clang::QualType pointee_type; /**< 指向类型 */
  uint64_t base;                /**< 基地址 */
  int64_t offset;               /**< 偏移字节数 */
  int64_t index;                /** 指针指向元素索引 */
  FixedAddress(const clang::Expr *parent, const clang::QualType &type,
               const clang::QualType &pointee_type, const clang::Expr *expr,
               uint64_t base, int64_t offset, int64_t index)
      : parent(parent),
        expr(expr),
        type(type),
        pointee_type(pointee_type),
        base(base),
        offset(offset),
        index(index) {}
};

/**
 * @brief 获取绝对地址的模板类
 * 继承并取代clang::RecursiveASTVisitor模板。
 */
template <typename T>
class FixedAddressVisitor : public clang::RecursiveASTVisitor<T> {
  std::vector<clang::Stmt *> stmt_stack;

 protected:
  clang::ASTContext &ast;

 public:
  /** fixed addresses */
  std::map<std::string, std::vector<FixedAddress>> fixed_addresses;
  FixedAddressVisitor(clang::ASTContext &ast) : ast(ast) {}
  bool dataTraverseStmtPre(clang::Stmt *s) {
    stmt_stack.push_back(s);
    return true;
  }

  bool dataTraverseStmtPost(clang::Stmt *s) {
    stmt_stack.pop_back();
    return true;
  }

  clang::Expr *check_constant_fixed_address(clang::Expr *expr, uint64_t &base,
                                            int64_t &offset) {
    llvm::APSInt result;
    if (clang::isa<clang::CastExpr>(expr) ||
        clang::isa<clang::ParenExpr>(expr)) {
      clang::Expr *sub_expr;
      if (clang::CastExpr *ce = clang::dyn_cast<clang::CastExpr>(expr)) {
        sub_expr = check_constant_fixed_address(ce->getSubExpr(), base, offset);
      } else if (clang::ParenExpr *pe =
                     clang::dyn_cast<clang::ParenExpr>(expr)) {
        sub_expr = check_constant_fixed_address(pe->getSubExpr(), base, offset);
      }
      if (sub_expr) return sub_expr;

#if LLVM_VERSION_MAJOR > 10
      if (!expr->isValueDependent() && expr->isIntegerConstantExpr(ast)) {
#if LLVM_VERSION_MAJOR < 16
        result = expr->getIntegerConstantExpr(ast).getValue();
#else
        result = expr->getIntegerConstantExpr(ast).value();
#endif
        base = result.getExtValue();
        offset = 0;
        return expr;
      }
#else
      if (expr->isIntegerConstantExpr(result, ast)) {
        base = result.getExtValue();
        offset = 0;
        return expr;
      }
#endif

      return NULL;
    }
    if (clang::BinaryOperator *op =
            clang::dyn_cast<clang::BinaryOperator>(expr)) {
      clang::BinaryOperator::Opcode opc = op->getOpcode();
      if (opc == clang::BO_Comma) {
        return check_constant_fixed_address(op->getRHS(), base, offset);
      } else if (opc == clang::BO_Assign) {
        return check_constant_fixed_address(op->getRHS(), base, offset);
      }
      if (opc == clang::BO_Add || opc == clang::BO_Sub) {
        clang::Expr *lhs =
            check_constant_fixed_address(op->getLHS(), base, offset);
        if (!lhs) return NULL;
        clang::Expr *rhs = op->getRHS();
#if LLVM_VERSION_MAJOR > 10
        if (!rhs->isValueDependent() && rhs->isIntegerConstantExpr(ast)) {
#if LLVM_VERSION_MAJOR < 16
          result = rhs->getIntegerConstantExpr(ast).getValue();
#else
          result = rhs->getIntegerConstantExpr(ast).value();
#endif
          if (opc == clang::BO_Add) {
            offset += result.getExtValue();
          } else {
            offset -= result.getExtValue();
          }
          if (offset < 0) {
            base = base + offset;
            offset = 0;
          }
        }
#else
        if (rhs->isIntegerConstantExpr(result, ast)) {
          if (opc == clang::BO_Add) {
            offset += result.getExtValue();
          } else {
            offset -= result.getExtValue();
          }
          if (offset < 0) {
            base = base + offset;
            offset = 0;
          }
        }
#endif
        return lhs;
      }
    }

#if LLVM_VERSION_MAJOR > 10
    if (!expr->isValueDependent() && expr->isIntegerConstantExpr(ast)) {
#if LLVM_VERSION_MAJOR < 16
      result = expr->getIntegerConstantExpr(ast).getValue();
#else
      result = expr->getIntegerConstantExpr(ast).value();
#endif
      base = result.getExtValue();
      offset = 0;
      return expr;
    }
#else
    if (expr->isIntegerConstantExpr(result, ast)) {
      base = result.getExtValue();
      offset = 0;
      return expr;
    }
#endif
    return NULL;
  }

  void add_fixed_address(std::vector<clang::Stmt *>::const_iterator it,
                         const clang::QualType &pointer_type, clang::Expr *expr,
                         uint64_t base = 0, int64_t offset = 0,
                         int64_t index = 0) {
    assert(pointer_type->isPointerType());
    clang::QualType pointee_type =
        pointer_type->getPointeeType().getCanonicalType().getUnqualifiedType();
    if (it != stmt_stack.begin() && clang::isa<clang::Expr>(*(it - 1))) {
      --it;
      if (clang::CastExpr *ce = clang::dyn_cast<clang::CastExpr>(*it)) {
        clang::QualType type = ce->getType();
        if (type->isPointerType()) {
          if (!pointee_type->isDependentType()) {
            const auto *RT = clang::dyn_cast<const clang::RecordType>(
                pointee_type.getTypePtr());
            if (!RT || RT->getDecl()->getDefinition()) {
              offset +=
                  index * ast.getTypeSizeInChars(pointee_type).getQuantity();
            }
            add_fixed_address(it, type, expr, base, offset, 0);
          }
          return;
        }
      } else if (clang::ArraySubscriptExpr *ase =
                     clang::dyn_cast<clang::ArraySubscriptExpr>(*it)) {
        // assert(ase->getType() == pointee_type);
        clang::Expr *idx = ase->getIdx();
        llvm::APSInt result;

#if LLVM_VERSION_MAJOR > 10
        if (!idx->isValueDependent() && idx->isIntegerConstantExpr(ast)) {
#if LLVM_VERSION_MAJOR < 16
          result = idx->getIntegerConstantExpr(ast).getValue();
#else
          result = idx->getIntegerConstantExpr(ast).value();
#endif
          index += result.getExtValue();
        }
#else
        if (idx->isIntegerConstantExpr(result, ast)) {
          index += result.getExtValue();
        }
#endif

        if (it != stmt_stack.begin()) {
          --it;
          if (clang::UnaryOperator *op =
                  clang::dyn_cast<clang::UnaryOperator>(*it)) {
            clang::UnaryOperator::Opcode opc = op->getOpcode();
            if (opc == clang::UO_AddrOf) {
              add_fixed_address(it, pointer_type, expr, base, offset, index);
              return;
            }
          }
          ++it;
        }
      } else if (clang::UnaryOperator *op =
                     clang::dyn_cast<clang::UnaryOperator>(*it)) {
        clang::UnaryOperator::Opcode opc = op->getOpcode();
        if (opc == clang::UO_Deref) {
          if (it != stmt_stack.begin()) {
            --it;
            if (clang::UnaryOperator *op =
                    clang::dyn_cast<clang::UnaryOperator>(*it)) {
              clang::UnaryOperator::Opcode opc = op->getOpcode();
              if (opc == clang::UO_AddrOf) {
                add_fixed_address(it, pointer_type, expr, base, offset, index);
                return;
              }
            }
            ++it;
          }
        }
      } else {
        clang::QualType type = clang::cast<clang::Expr>(*it)->getType();
        if (type->isPointerType()) {
          if (type->getPointeeType().getCanonicalType().getUnqualifiedType() ==
              pointee_type) {
            if (clang::BinaryOperator *op =
                    clang::dyn_cast<clang::BinaryOperator>(*it)) {
              clang::BinaryOperator::Opcode opc = op->getOpcode();
              if (opc == clang::BO_Add || opc == clang::BO_Sub) {
                clang::Expr *rhs = op->getRHS();
                llvm::APSInt result;
#if LLVM_VERSION_MAJOR > 10
                if (!rhs->isValueDependent() &&
                    rhs->isIntegerConstantExpr(ast)) {
#if LLVM_VERSION_MAJOR < 16
                  result = rhs->getIntegerConstantExpr(ast).getValue();
#else
                  result = rhs->getIntegerConstantExpr(ast).value();
#endif
                  if (opc == clang::BO_Sub) {
                    index -= result.getExtValue();
                  } else {
                    index += result.getExtValue();
                  }
                  if (index < 0) {
                    offset +=
                        index *
                        ast.getTypeSizeInChars(pointee_type).getQuantity();
                    index = 0;
                    if (offset < 0) {
                      base = base + offset;
                      offset = 0;
                    }
                  }
                }
#else
                if (rhs->isIntegerConstantExpr(result, ast)) {
                  if (opc == clang::BO_Sub) {
                    index -= result.getExtValue();
                  } else {
                    index += result.getExtValue();
                  }
                  if (index < 0) {
                    offset +=
                        index *
                        ast.getTypeSizeInChars(pointee_type).getQuantity();
                    index = 0;
                    if (offset < 0) {
                      base = base + offset;
                      offset = 0;
                    }
                  }
                }
#endif
              }
            }
            add_fixed_address(it, pointer_type, expr, base, offset, index);
            return;
          }
        }
        ++it;
      }
    }
    clang::Expr *parent = clang::dyn_cast<clang::Expr>(*it);
    assert(parent != NULL);
    std::string parent_string;
    {
      llvm::raw_string_ostream os(parent_string);
      parent->printPretty(os, nullptr,
                          clang::PrintingPolicy(ast.getLangOpts()));
      os.flush();
    }
    auto &addresses = fixed_addresses[parent_string];
    // Until clang 9.0, an expression can be visited more than once.
    if (addresses.empty() || addresses.back().parent != parent) {
      addresses.emplace_back(parent, pointer_type, pointee_type, expr, base,
                             offset, index);
    }
  }

  void check_integer_to_pointer_cast(clang::Expr *expr, clang::Expr *sub_expr,
                                     bool allow_ptr2func = false) {
    clang::QualType dest_type = expr->getType();
    if (!dest_type->isPointerType() ||
        (dest_type->isFunctionPointerType() && !allow_ptr2func))
      return;
    // Fixme: is NULL value dependent or not?
    if (expr->isNullPointerConstant(ast,
                                    clang::Expr::NPC_ValueDependentIsNotNull))
      return;
    clang::QualType src_type = sub_expr->getType();
    if (!src_type->isIntegralType(ast) || src_type->isBooleanType() ||
        src_type->isEnumeralType()) {
      return;
    }
    llvm::APSInt result;
#if LLVM_VERSION_MAJOR > 10
    if (!sub_expr->isValueDependent() && sub_expr->isIntegerConstantExpr(ast)) {
#if LLVM_VERSION_MAJOR < 16
      result = sub_expr->getIntegerConstantExpr(ast).getValue();
#else
      result = sub_expr->getIntegerConstantExpr(ast).value();
#endif
      // ignore the maximum address, eg: ~0、-1
      if (result.getExtValue() != -1) {
        add_fixed_address(stmt_stack.end(), dest_type, sub_expr,
                          result.getExtValue(), 0);
      }
      return;
    }
#else
    if (sub_expr->isIntegerConstantExpr(result, ast)) {
      // ignore the maximum address, eg: ~0、-1
      if (result.getExtValue() != -1) {
        add_fixed_address(stmt_stack.end(), dest_type, sub_expr,
                          result.getExtValue(), 0);
      }
      return;
    }
#endif

    uint64_t base;
    int64_t offset;
    if (clang::ConditionalOperator *op =
            clang::dyn_cast<clang::ConditionalOperator>(sub_expr)) {
      // 三元条件表达式可能产生两个结果
      if (clang::Expr *base_expr =
              check_constant_fixed_address(op->getLHS(), base, offset)) {
        add_fixed_address(stmt_stack.end(), dest_type, op->getLHS(), base,
                          offset);
      } else {
        // 隐式绝对地址
        add_fixed_address(stmt_stack.end(), dest_type, op->getLHS());
      }
      if (clang::Expr *base_expr =
              check_constant_fixed_address(op->getRHS(), base, offset)) {
        add_fixed_address(stmt_stack.end(), dest_type, op->getRHS(), base,
                          offset);
      } else {
        // 隐式绝对地址
        add_fixed_address(stmt_stack.end(), dest_type, op->getRHS());
      }
    } else if (clang::Expr *base_expr =
                   check_constant_fixed_address(sub_expr, base, offset)) {
      add_fixed_address(stmt_stack.end(), dest_type, sub_expr, base, offset);
    } else {
      // 隐式绝对地址
      add_fixed_address(stmt_stack.end(), dest_type, sub_expr);
    }
  }

  bool VisitImplicitCastExpr(clang::ImplicitCastExpr *ice) {
    // 检查隐含类型转换到指针
    check_integer_to_pointer_cast(ice, ice->getSubExpr(), false);
    return true;
  }
  bool VisitCStyleCastExpr(clang::CStyleCastExpr *cce) {
    // 检查强制类型转换到指针
    if (!cce) return true;
    if (cce->getCastKind() == clang::CK_NullToPointer) return true;
    if (auto sub_expr = cce->getSubExpr()) {
      check_integer_to_pointer_cast(cce, sub_expr,
                                    sub_expr->getType()->isIntegerType());
    }
    return true;
  }
};

}  // namespace share
}  // namespace rocket

#endif /* FIXED_ADDRESS_H */
